from drf_yasg.inspectors import SwaggerAutoSchema

URL_GROUPS = {
    "auth": '权限',
    'sms': '手机验证',
    'user': '用户',
    'token': '权限',
    "labs": '实验室',
    'message': "消息",
    'credit': '信用分'
}


def get_group(url):
    return URL_GROUPS.get(url, url)


class CustomSwaggerAutoSchema(SwaggerAutoSchema):

    def get_tags(self, operation_keys=None):
        tags = super().get_tags(operation_keys)
        tags = list(map(get_group, tags))
        return tags

    def get_summary_and_description(self):
        description = self.overrides.get('operation_description', None)
        summary = self.overrides.get('operation_summary', None)
        if description is None:
            description = self._sch.get_description(self.path, self.method) or ''
            description = description.strip().replace('\r', '')

            if description and (summary is None):
                summary_max_len, sections = 120, description.split('\n\n', 1)
                if len(sections) == 2 and len(sections[0].strip()) < summary_max_len:
                    _, description = sections
                    summary = description.strip()

        return description, summary
